#include "StdAfx.h"
#include "SocketHelper.h"
#include <process.h>

#ifdef CREATE_CONSOLE
#include <iostream>
extern HANDLE m_console;
#endif

CSocketHelper::CSocketHelper(VOID) : m_socketRelated(0), m_isCacheDataAvailable(CDS_UNDEFINED)
{
}

CSocketHelper::CSocketHelper(SOCKET s)
{
	m_socketRelated = s;
	m_isCacheDataAvailable = CDS_UNDEFINED;
}

CSocketHelper::CSocketHelper(CSocketHelper &helper)
{
	this->m_socketRelated = helper.m_socketRelated;
	this->m_isCacheDataAvailable = helper.m_isCacheDataAvailable;
}

CSocketHelper::~CSocketHelper(void)
{
}


BOOL CSocketHelper::IsCacheDataAvailable(void)
{
// 	CHAR buffer[100] = {0};
// 	sprintf(buffer, "Related socket = %x, data available? %d\n", m_socketRelated, m_isCacheDataAvailable);
// 	WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
	if(m_isCacheDataAvailable == CDS_UNDEFINED)
	{
		// Ask maincpp
		m_isCacheDataAvailable = CheckAvailableCacheData(m_socketRelated);
	}
	if(m_isCacheDataAvailable == CDS_AVAILABLE || m_isCacheDataAvailable == CDS_PENDING)
	{
		return TRUE;
	}
	else
	{
		return FALSE;
	}
}

BOOL CSocketHelper::IsFollowingCacheDataAvailable(void)
{
	if(m_isCacheDataAvailable == CDS_UNDEFINED)
	{
		// Ask maincpp
		m_isCacheDataAvailable = CheckAvailableCacheData(m_socketRelated);
	}
	if(m_isCacheDataAvailable == CDS_NOTAVAILABLE)
	{
		return FALSE;
	}
	else
	{
		INT nNewState = CheckAvailableCacheData(m_socketRelated);
		if(nNewState == CDS_AVAILABLE)
		{
			return TRUE;
		}
		else
		{
			// Pending
			return FALSE;
		}
	}
}


VOID CSocketHelper::SetRelatedSocket(SOCKET socket)
{
	m_socketRelated = socket;
}


INT CSocketHelper::GetCacheData(PUCHAR lpData, int nLength)
{
	int ret = GetAvailableCacheData(m_socketRelated, lpData, nLength);
	return ret;
}

INT CSocketHelper::GetCacheData(LPWSABUF lpBuffers, DWORD nBufferCount)
{
	INT ret = 0;
	int nLength = 0;
	for(int i = 0; i < nBufferCount; ++i)
	{
		nLength += lpBuffers[i].len;
	}

	PUCHAR lpData = new UCHAR[nLength];
	int nDataSize = GetAvailableCacheData(m_socketRelated, lpData, nLength);
	
	if(nDataSize == 0)
	{
		// No more data available
		// TODO: what to do here?
		ret = 0;
	}
	else if(nDataSize == -1)
	{
		// We should wait for more time!
		WSASetLastError(WSA_IO_PENDING);
		ret = SOCKET_ERROR;
	}
	else
	{
		// Fill in the buffers
		DWORD nBytesWritten = 0;
		for(int i = 0; i < nBufferCount && nBytesWritten < nDataSize; ++i)
		{
			DWORD nBytesToWrite = min(lpBuffers[i].len, nDataSize - nBytesWritten);
			memcpy_s(lpBuffers[i].buf, lpBuffers[i].len, lpData + nBytesWritten, nBytesToWrite);
			nBytesWritten += nBytesToWrite;
		}
		ret = nBytesWritten;
	}

	delete[] lpData;

	return ret;
}

CWSABufferInfo* CSocketHelper::RegisterWSAEvent(WSAEVENT uEvent, LPWSABUF lpBuffers, DWORD dwBufferCount, LPDWORD lpNumberOfBytesRecvd, WinsockHookType originalHookType)
{
	if(mapEvent2WSABuffer[uEvent] != NULL)
	{
		CWSABufferInfo* p = mapEvent2WSABuffer[uEvent];
		delete p;
	}

	CWSABufferInfo* pInfo = new CWSABufferInfo;
	pInfo->dwBufferCount = dwBufferCount;
	pInfo->uNumberOfBytesRecvd = *lpNumberOfBytesRecvd;
	pInfo->lpBuffers = lpBuffers;
	pInfo->originalHookType = originalHookType;
	mapEvent2WSABuffer[uEvent] = pInfo;

	return pInfo;
}

BOOL CSocketHelper::RemoveWSAEvent(WSAEVENT uEvent)
{
	mapEvent2WSABuffer[uEvent] = NULL;
	return TRUE;
}

CWSABufferInfo* CSocketHelper::GetWSABufferFromEvent(WSAEVENT uEvent)
{
	// TODO: locks!
	return mapEvent2WSABuffer[uEvent];
}

VOID CSocketHelper::StartCacheStatusGuardingThread(HANDLE hEvent)
{
	if(m_bIsCSGTStarted == FALSE)
	{
		m_bIsCSGTStarted = TRUE;
		// Start a new thread to reset the signal
		CCacheStatusGuardInfo *pCSGI = new CCacheStatusGuardInfo;
		pCSGI->pHelper = this;
		pCSGI->hEvent = hEvent;

#ifdef CREATE_CONSOLE
		CHAR buffer[100] = {0};
		sprintf(buffer, "Start thread!\n");
		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
		_beginthread(threadCacheStatusGuard, 0, pCSGI);
	}
}

VOID CSocketHelper::threadCacheStatusGuard(PVOID param)
{
	CCacheStatusGuardInfo *pInfo = (CCacheStatusGuardInfo*)param;

#ifdef CREATE_CONSOLE
	CHAR buffer[100] = {0};
	sprintf(buffer, "Thread func executed!\n");
	WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

// 	while(true)
// 	{
		while(!pInfo->pHelper->IsFollowingCacheDataAvailable())
		{
#ifdef CREATE_CONSOLE
			CHAR buffer[100] = {0};
			sprintf(buffer, "CacheStatusGuard - The following data is not available. Keep on looping.\n");
			WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
			Sleep(300);
		}

#ifdef CREATE_CONSOLE
		sprintf(buffer, "CacheStatusGuard - The following data is available. Set signal.");
		WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif

		WSASetEvent(pInfo->hEvent);

		pInfo->pHelper->m_bIsCSGTStarted = FALSE;
// 		Sleep(300);
// 	}
	

	delete pInfo;

#ifdef CREATE_CONSOLE
	sprintf(buffer, "Thread func exited!\n");
	WriteConsoleA(m_console, buffer, strlen(buffer), NULL, NULL);
#endif
}